import json
import math
import time
# 多线程
from collections import deque
from threading import Thread

import cv2
import gym
import numpy as np
import win32con
import win32gui
import win32ui
from pykeyboard import PyKeyboard
from redis import Redis


# from pymouse import PyMouse


class HKEnv(gym.Env):
    def __init__(self):
        self.r = Redis(host='127.0.0.1', port=6379, db=1)
        self.hwnd_id = []
        self.k = PyKeyboard()
        self.initailize()
        win32gui.EnumWindows(self.get_all_hwnd, 0)
        self.frames = deque([], maxlen=4)
        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)
        # 创建dc, 截取桌面
        self.hdesktop = win32gui.GetDesktopWindow()
        # 获取窗口大小
        l, t, r, b = win32gui.GetWindowRect(self.hwnd_id[0])
        self.width = r - l
        self.height = b - t

        # 创建设备描述表
        self.desktop_dc = win32gui.GetWindowDC(self.hdesktop)
        # 创建dc对象
        self.mfc_dc = win32ui.CreateDCFromHandle(self.desktop_dc)
        # 创建内存设备描述表
        self.save_dc = self.mfc_dc.CreateCompatibleDC()
        # 创建位图对象
        self.save_bit_map = win32ui.CreateBitmap()
        # 截图
        self.save_bit_map.CreateCompatibleBitmap(self.mfc_dc, self.width, self.height)

        # self.action_list = ['攻击', '左+攻', '右+攻', '上+攻', '下+攻', '左+跳', '右+跳', '左+跳+攻', '右+跳+攻', '波', '吼', '砸', '冲']
        # self.action_dict = {
        #     0: '攻击',
        #     1: '左+攻',
        #     2: '右+攻',
        #     3: '上+攻',
        #     4: '下+攻',
        #     5: '左+跳',
        #     6: '右+跳',
        #     7: '左+跳+攻',
        #     8: '右+跳+攻',
        #     9: '波',
        #     10: '吼',
        #     11: '砸',
        #     12: '冲'
        # }
        self._max_episode_steps = 10000
        self.observation_space = gym.spaces.Box(low=0, high=100, shape=(1, 8), dtype=np.float32)
        # self.observation_space = gym.spaces.Discrete(8)
        self.action_space = gym.spaces.Discrete(12)
        self.std = 10
        self.mean = -1
        self.steps = 2

    def angle(self, x, y):
        a = np.arctan2(y, x)
        if a < 0:
            a = a + 2 * np.pi
        return a

    def calculate_mean_std(self, new_history, old_mean, old_std):
        # sample_count：旧样本总数
        # new_sample_num：新样本总数
        # new_history：新的样本
        # old_mean：旧均值
        # old_var：旧方差
        new_mean = (self.steps * old_mean + new_history)
        new_mean /= (1 + self.steps)
        new_std = np.sum(np.square(new_history - new_mean), axis=0)
        new_std = (self.steps * np.square(old_std) + new_std)
        new_std /= (1 + self.steps)
        new_std = np.sqrt(new_std)
        return new_mean, new_std

    def initailize(self):
        pass

    def get_all_hwnd(self, hwnd, mouse):
        if (win32gui.IsWindow(hwnd) and
            win32gui.IsWindowEnabled(hwnd) and
            win32gui.IsWindowVisible(hwnd)) and win32gui.GetWindowText(hwnd).__contains__('Hollow Knight'):
            self.hwnd_id.append(hwnd)

    def restart_game(self):
        win32gui.EnumWindows(self.get_all_hwnd, 0)
        time.sleep(7)
        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        time.sleep(1)
        if int(self.r.get('done').decode('utf-8')) == 1:
            time.sleep(6)
        time.sleep(8)
        self.k.release_key(self.k.left_key)
        self.k.release_key(self.k.right_key)
        self.k.release_key(self.k.up_key)
        self.k.release_key(self.k.down_key)
        self.k.release_key("Z")
        self.k.release_key("X")
        self.k.release_key("C")
        self.k.release_key("F")
        time.sleep(1)
        self.k.tap_key(self.k.left_key)
        time.sleep(4)
        self.k.tap_key(self.k.up_key)
        time.sleep(2)
        self.k.tap_key("z")

    def reset(self, *args, **kwargs):
        start = None
        obs = None
        while start is None:
            self.min_x = 10000
            self.max_x = 0
            self.face = 1
            self._last_enemy_hp = 0
            self._last_knight_hp = 0
            self.can_shadow_dash = True
            self.shadow_dashing_time = 0
            self.s = time.time()
            self.e = time.time()
            t = Thread(target=self.restart_game)
            t.start()
            start = self.r.brpop(['start'], timeout=30)
            time.sleep(0.1)
            state_json = self.r.get('hollow_obs')
            # binary = self.get_binary(save=True, no=0)
            obs = self.state_to_obs(state_json)
            self._last_enemy_hp = self.enemy_hp
            self._last_knight_hp = self.knight_hp
        print("reset done")
        return obs

    def do_action(self, action):
        # self.action_dict = {0: '攻击', 1: '上+攻', 2: '下+攻', 3: '跳', 4: '松+跳(双跳)',
        #                     5: '波', 6: '吼', 7: '砸', 8: '冲', 9: '左', 10: '右', 11: '松(无动作)'}
        L = self.k.left_key
        R = self.k.right_key
        U = self.k.up_key
        D = self.k.down_key
        if action == 0:
            self.k.tap_key('X')
        if action == 1:
            self.k.release_key(D)
            self.k.press_key(U)
            time.sleep(0.01)
            self.k.tap_key('X')
            time.sleep(0.01)
            self.k.release_key(U)
        if action == 2:
            self.k.release_key(U)
            self.k.press_key(D)
            time.sleep(0.01)
            self.k.tap_key('X')
            time.sleep(0.01)
            self.k.release_key(D)
        if action == 3:
            self.k.press_key('Z')
        if action == 4:
            self.k.release_key('Z')
            time.sleep(0.01)
            self.k.press_key('Z')
        if action == 5:
            self.k.tap_key('F')
        if action == 6:
            self.k.release_key(D)
            self.k.press_key(U)
            time.sleep(0.01)
            self.k.tap_key('F')
            time.sleep(0.01)
            self.k.release_key(U)
        if action == 7:
            self.k.release_key(U)
            self.k.press_key(D)
            time.sleep(0.01)
            self.k.tap_key('F')
            time.sleep(0.01)
            self.k.release_key(D)
        if action == 8:
            self.k.tap_key("C")
        if action == 9:
            self.face = 0.5
            self.k.release_key(R)
            self.k.press_key(L)
        if action == 10:
            self.face = 1
            self.k.release_key(L)
            self.k.press_key(R)
        if action == 11:
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.release_key('Z')

    def get_binary(self, save=False, no=0):
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)

        self.save_dc.SelectObject(self.save_bit_map)
        # 截取从左上角(50, 50)长宽为(width, height)的图片
        # self.save_dc.BitBlt((0, 0), (self.width - 21, self.height - 41), self.mfc_dc, (0, 0), win32con.SRCCOPY)
        self.save_dc.BitBlt((0, 0), (self.width - 300, self.height - 160), self.mfc_dc, (150, 50), win32con.SRCCOPY)

        # 保存位图文件
        signedIntsArray = self.save_bit_map.GetBitmapBits(True)
        img = np.frombuffer(signedIntsArray, dtype='uint8')
        img.shape = (self.height, self.width, 4)
        # 删除多余部分
        img = img[0:self.height - 160, 0:self.width - 300]
        # 缩放
        img = cv2.resize(img, (300, 300))
        if save:
            cv2.imwrite("imgs/{}.png".format(no), img)
        # 保存
        # cv2.imwrite('test.png', img)
        # 转换为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if save:
            cv2.imwrite("imgs/gray_{}.png".format(no), gray)
        # 压缩到0-1
        gray = gray / 255.0
        # 缩放
        # gray = cv2.resize(gray, (511, 511))
        # 保存
        # cv2.imwrite('test_gray.png', gray)
        # 二值化
        # ret, binary = cv2.threshold(gray, 127, 1, cv2.THRESH_BINARY)
        # if save:
        #     ret, binary_255 = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        #     cv2.imwrite("imgs/binary_{}.png".format(no), binary_255)
        # 转成numpy数组
        binary = np.array(gray)
        return binary

    def step(self, action):
        if action == 8:
            if self.can_shadow_dash:
                self.can_shadow_dash = False
                self.shadow_dashing_time = 0
        reward = 0
        if action == 11:
            reward = -1
        if abs(self.enemy_x - self.knight_x) > 2:
            if self.enemy_x > self.knight_x:
                # 向左走
                if action == 9:
                    reward -= 1
                # 左冲
                if self.face == 0.5 and action == 8:
                    reward -= 1
            else:
                # 向右走
                if action == 10:
                    reward -= 1
                # 右冲
                if self.face == 1 and action == 8:
                    reward -= 1

        if self.knight_mp < 33 and action in [5, 6, 7]:
            reward -= 1
        if action in [0, 1, 2, 5, 6, 7]:
            reward -= 1
        if self.is_attacking and action in [0, 1, 2, 5, 6, 7]:
            reward -= 1

        if self.distance < 3.5 and not self.is_attacking:
            # boss在上方,也就是弧度在 3pi/8 ~ 5pi/8
            if self.angle_kb > 3 * math.pi / 8 and self.angle_kb < 5 * math.pi / 8:
                if action == 1:
                    reward = self.distance * 10
            # boss在下方,也就是弧度在 11pi/8 ~ 13pi/8
            elif self.angle_kb > 7 * math.pi / 8 and self.angle_kb < 9 * math.pi / 8:
                if action == 2:
                    reward = self.distance * 10
            # boss在左方,也就是弧度在 7pi/8 ~ 9pi/8
            elif self.angle_kb > 7 * math.pi / 8 and self.angle_kb < 9 * math.pi / 8:
                if self.face == 0.5:
                    if action == 0:
                        reward = self.distance * 10
            # boss在右方,也就是弧度小于 1pi/8 或者大于 15pi/8
            elif self.angle_kb < math.pi / 8 or self.angle_kb > 15 * math.pi / 8:
                if self.face == 1:
                    if action == 0:
                        reward = self.distance * 10

        start = time.time()
        self.do_action(action)
        max_time = 0.03
        while True:
            end = time.time()
            state_json = self.r.get('hollow_obs')
            self.state_to_obs(state_json)
            if str(self.r.get('done').decode('utf-8')) == "1":
                print('kill')
                break
            if self.knight_spell or self.dashing or end - start < max_time:
                time.sleep(0.001)
                continue
            if end - start < max_time + 0.01:
                continue
            else:
                break
        duration = time.time() - start
        if self.knight_x - self.min_x < 0.05:
            reward -= 1
        if self.max_x - self.knight_x < 0.05:
            reward -= 1
        if not self.can_shadow_dash:
            self.shadow_dashing_time += duration
        if self.shadow_dashing_time >= 1.7:
            self.can_shadow_dash = True
        state_json = self.r.get('hollow_obs')
        obs = self.state_to_obs(state_json)
        # distance_diff = distance_before - self.distance
        # if action not in [0, 1, 2, 5, 6, 7]:
        #     if self.distance > 5:
        #         if distance_diff > 0:
        #             reward += 1 + distance_diff
        #         else:
        #             reward += distance_diff
        # if self.distance < 3.5 and distance_before >= self.distance:
        #     reward -= 5
        if self.distance > 5:
            reward -= 1
        self.done = False
        self.key_point = False
        if self._last_enemy_hp > self.enemy_hp:
            # if self.distance > 3:
            # self.key_point = True
            if action in [0, 1, 2]:
                reward = self.distance * 10
            elif action in [5, 6]:
                reward = (self.distance + 3) * 10
            elif action == 7:
                reward = self.distance * 10
        if self.distance < 2:
            reward -= 10
        if self._last_knight_hp > self.knight_hp:
            reward = -100
            self.key_point = True
        if self.knight_hp <= 0 or (self._last_knight_hp < 5 and self.knight_hp > 6):
            self.done = True
            reward = -100
            self.key_point = True
        if str(self.r.get('done').decode('utf-8')) == "1":
            self.done = True
            # if self.distance > 3.5:
            #     self.key_point = True
            reward = self.distance * 10
        self._last_enemy_hp = self.enemy_hp
        self._last_knight_hp = self.knight_hp
        if self.done:
            # self.get_binary(save=True, no=1)
            self.k.release_key(self.k.left_key)
            self.k.release_key(self.k.right_key)
            self.k.release_key(self.k.up_key)
            self.k.release_key(self.k.down_key)
            self.k.release_key("Z")
            self.k.release_key("X")
            self.k.release_key("C")
            self.k.release_key("F")
        # 缩放reward
        self.mean, self.std = self.calculate_mean_std(reward, self.mean, self.std)
        reward = reward / self.std
        self.steps += 1
        return obs, reward, self.done, self.key_point

    def state_to_obs(self, state):
        state_dict = json.loads(state)
        # 是否黑冲
        self.shadow_dashing = True if state_dict['shadow_dashing'] == '1' else False

        # 是否冲刺
        self.dashing = True if state_dict['dashing'] == '1' else False

        # 是否在攻击
        self.is_attacking = True if state_dict['is_attacking'] == '1' else False

        # 是否在释放技能
        self.knight_spell = True if state_dict['knight_spell'] != "Inactive" else False

        # 是否在地面
        self.on_ground = True if state_dict['onGround'] == '1' else False

        # 血量
        self.knight_hp = int(state_dict['knight_hp'])
        self.enemy_hp = int(state_dict['enemy_hp'])

        # knight 的位置
        self.knight_x = float(state_dict['knight_x'])
        self.knight_y = float(state_dict['knight_y'])

        # knight的灵魂量
        self.knight_mp = float(state_dict['knight_mp'])

        # boss 的位置
        self.enemy_x = float(state_dict['enemy_x'])
        self.enemy_y = float(state_dict['enemy_y'])

        # 距离, 用于判断是否在攻击范围内, 平方和开方,
        self.distance = math.sqrt((self.knight_x - self.enemy_x) ** 2 + (self.knight_y - self.enemy_y) ** 2)

        # 冲刺是否冷却
        self.is_sprint = True

        # 相对位置
        self.relative_x = self.knight_x - self.enemy_x
        self.relative_y = self.knight_y - self.enemy_y

        self.angle_kb = self.angle(-self.relative_x, -self.relative_y)

        if self.knight_x < self.min_x:
            self.min_x = self.knight_x
        if self.knight_x > self.max_x:
            self.max_x = self.knight_x

        obs = [self.knight_x / 100, self.knight_y / 100, self.knight_mp / 100, self.face, self.is_attacking, self.can_shadow_dash, self.angle_kb, self.distance]
        return np.array([obs], dtype=np.float32)


if __name__ == '__main__':
    env = HKEnv()
    # env.restart_game()
    win32gui.EnumWindows(env.get_all_hwnd, 0)
    handle = env.hwnd_id[0]
    # 激活窗口
    win32gui.SetForegroundWindow(handle)
    time.sleep(3)
    k = PyKeyboard()
    k.press_key("Z")
    time.sleep(0.2)
    k.release_key("Z")
    time.sleep(0.1)
    k.press_key("Z")
    time.sleep(1)
    k.release_key('Z')

    obs = env.reset()
    print(obs)
    for i in range(100):
        action = np.random.randint(0, 12)
        obs, reward, done, info = env.step(1, 1)
        count = 0
        print(obs, reward, done, info)
